import json
import numpy as np
import matplotlib.pyplot as plt

# 假设这是你的JSON列表
with open("generative_pm/results_Llama3_pairwise_reference.json", 'r') as file:
    # 加载JSON内容到字典
    data = json.load(file)

chosen_score_list = data["chosen_score_list"]

# 将列表转换为numpy数组方便计算
scores = np.array(chosen_score_list)

# 计算每个分数距离0.5的绝对值
distances = np.abs(scores - 0.5)

# 定义bins的边界，这里以0.1为一个bin的宽度，限制在0到0.5
bins = np.arange(0, 0.6, 0.1)

# 计算每个分数属于哪个bin
bin_indices = np.digitize(distances, bins, right=True) - 1

# 初始化一个字典来存储每个bin的结果
bin_results = {}

# 遍历每个bin
for i in range(len(bins) - 1):
    # 找出当前bin中的所有分数
    bin_scores = scores[bin_indices == i]
    
    # 计算大于0.5的数量
    greater_than_half = np.sum(bin_scores > 0.5)
    
    # 计算比例
    proportion_greater = greater_than_half / len(bin_scores) if bin_scores.size > 0 else 0
    
    # 存储结果，包括大于0.5的比例和该bin的样本数
    bin_results[(bins[i], bins[i+1])] = {
        'greater_than_half': proportion_greater,
        'sample_size': len(bin_scores)
    }

# 准备数据用于绘图
bin_ranges = [(f"{k[0]+0.5:.1f}-{k[1]+0.5:.1f}", v['greater_than_half'], v['sample_size']) for k, v in bin_results.items() if k[0] <= 0.5]

# 绘制柱状图
fig, ax = plt.subplots()
bar_width = 0.7
index = np.arange(len(bin_ranges))

# 绘制大于0.5的柱状图
ax.bar(index, [x[1] for x in bin_ranges], bar_width, label='Greater than 0.5')

# 添加标签和标题
ax.set_xlabel('Preference Strength')
ax.set_ylabel('Accuracy')
ax.set_title('Ultra Feedback')
ax.set_xticks(index)
ax.set_xticklabels([f"{r[0]}" for r in bin_ranges])
ax.legend()
ax.set_ylim(0.4, 0.8)

# 在每个柱子上方显示样本数量
for i, (_, proportion, sample_size) in enumerate(bin_ranges):
    ax.text(i, proportion + 0.02, f'N={sample_size}', ha='center', va='bottom')

# 显示图形
plt.tight_layout()
# 保存图形到文件
plt.savefig('generative_pm/proportion_bins-Llama3_pairwise_reference.png')
plt.show()
